package dslab.util;

import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Path;
import java.security.*;
import java.util.Base64;

/**
 * utility class that contains all sorts of static methods used to facilitate secure and encrypted communication
 * includes base64 helpers, hmac helpers, RSA helpers and AES helpers
 */
public class CryptoUtil {

   /**
     * generates a new AES secret key
     * @return AES Secret Key
     * @throws NoSuchAlgorithmException should not happen, but may if library has problems
     */
    public static SecretKey generateSecret() throws NoSuchAlgorithmException {
        KeyGenerator keyGenerator;
        try {
            keyGenerator = KeyGenerator.getInstance("AES");
        } catch (NoSuchAlgorithmException exception) {
            // Should not happen
            throw exception;
        }
        keyGenerator.init(256);
        return keyGenerator.generateKey();
    }

    /**
     * creates an array of secure random numbers of specified size
     * @param size size of returned Array
     * @return a byte array of secure random numbers
     */
    public static byte[] generateRandom(int size) {
        byte[] rand = new byte[size];
        new SecureRandom().nextBytes(rand);
        return rand;
    }

    /**
     * creates a Pair object containing a secret key (AES) and an initialisation vector (as a byte array)
     * @return Pair containing a secret key and a byte array of size 16
     * @throws NoSuchAlgorithmException - Should never happen, could in case of a library failure
     */
    public static Pair<SecretKey, byte[]> generateNewKeydata() throws NoSuchAlgorithmException {
        return new Pair<>(generateSecret(), generateRandom(16));
    }

    /**
     * recovers an AES secret key from a String
     * @param key String containing an encoded secret key
     * @return a secret key recovered from the specified String
     */
    public static SecretKey decodeSecret(String key) {
        return new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), 0, key.length(), "AES");
    }

    /**
     * recovers an AES secret key from a byte Array
     * @param key String containing an encoded secret key
     * @return a secret key recovered from the specified byte Array
     */
    public static SecretKey decodeSecret(byte[] key) {
        return new SecretKeySpec(key, 0, key.length, "AES");
    }

    /**
     * encodes the specified String in Base64 and returns it as a String
     * @param message String to be encoded
     * @return a String encoded in Base6
     */
    public static String encodeB64(String message) {
        return Base64.getEncoder().encodeToString(message.getBytes(StandardCharsets.UTF_8));
    }

    /**
     * encodes the specified byte Array in Base64 and returns it as a String
     * @param message Byte Array to be encoded
     * @return a String encoded in Base64
     */
    public static String encodeB64(byte[] message) {
        return Base64.getEncoder().encodeToString(message);
    }

    /**
     * decodes the specified String from Base64 and returns it as a byte Array
     * @param message String to be decoded
     * @return A decoded byte Array
     */
    public static byte[] decodeB64(String message) {
        return Base64.getDecoder().decode(message);
    }

    /**
     * decodes the specified byte Array from Base64 and returns it as a byte Array
     * @param message String to be decoded
     * @return a decoded byte Array
     */
    public static byte[] decodeB64(byte[] message) {
        return Base64.getDecoder().decode(message);
    }

    /**
     * reads the private key from a specified file encoded in the .der format
     * @param location location of a file containing a private key
     * @return a private Key read from a file
     * @throws IOException if an IO error occurs
     */
    public static PrivateKey getPrivateKey(Path location) throws IOException{
        return Keys.readPrivateKey(location.toFile());
    }

    /**
     * reads the public key from a specified file encoded in the .der format
     * @param location location of a file containing a public key
     * @return a public Key read from a file
     * @throws IOException if an IO error occurs
     */
    public static PublicKey getPublicKey(Path location) throws IOException {
        return Keys.readPublicKey(location.toFile());
    }

    /**
     * reads the secret key from a specified file
     * @param location location of a file containing a secret key
     * @return a secret key read from a file
     * @throws IOException if an IO error occurs
     */
    public static SecretKey getSecretKey(Path location) throws IOException {
        return Keys.readSecretKey(location.toFile());
    }

    /**
     * creates a Cipher object for encrypting in RSA (asymmetric)
     * @param key key used for encrypting
     * @return a cipher initialised for encrypting
     */
    public static Cipher createRSAEncryptCipher(PublicKey key) {
        try {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.ENCRYPT_MODE, key);
            return cipher;
        } catch (Exception e) {
            System.err.println(e.getMessage());
            return null;
        }
    }

    /**
     * creates a Cipher object for decrypting in RSA (asymmetric)
     * @param key key used for decrypting
     * @return a cipher initialised for decrypting
     */
    public static Cipher createRSADecryptCipher(PrivateKey key) {
        try {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.DECRYPT_MODE, key);
            return cipher;
        } catch (Exception e) {
            System.err.println(e.getMessage());
            return null;
        }
    }

    /**
     * encrypts a specified String with a specified cipher and encodes the encrypted String in Base64
     * @param message Message to be encrypted
     * @param cipher Cipher to be used for encryption
     * @return an encrypted String in Base64
     * @throws IllegalBlockSizeException if the provided String does not match the block size of the cipher in case a block-cipher is used
     * @throws BadPaddingException if the input String has incorrect padding for the used cipher
     */
    public static String encryptMessage(String message, Cipher cipher) throws IllegalBlockSizeException, BadPaddingException {
        return new String(Base64.getEncoder().encode(cipher.doFinal(message.getBytes(StandardCharsets.UTF_8))));
    }

    /**
     * decrypts a specified String in Base64 with a specified cipher
     * @param message String in Base64 to be decrypted
     * @param cipher Cipher to be used for decryption
     * @return a decrypted String
     * @throws IllegalBlockSizeException if the provided String does not match the block size of the cipher in case a block-cipher is used
     * @throws BadPaddingException if the input String has incorrect padding for the used cipher
     */
    public static String decryptMessage(String message, Cipher cipher) throws IllegalBlockSizeException, BadPaddingException {
        return new String(cipher.doFinal(Base64.getDecoder().decode(message.getBytes(StandardCharsets.UTF_8))));
    }

    /**
     * hashes a Message object using the HmacSHA256 algorithm with the provided secret key
     * @param message Message to be hashed
     * @param key Key to be used for hashing
     * @return a string containing the hash of the specified String encoded in Base64
     * @throws NoSuchAlgorithmException should never happen, could in case of a library failure
     * @throws InvalidKeyException if an invalid key is provided
     */
    public static String hashMessage(Message message, SecretKey key) throws NoSuchAlgorithmException, InvalidKeyException {
        return encodeB64(hashMessageArr(message, key));
    }

    /**
     * verifies a HmacSHA256 hash with the provided secret key
     * @param message Message object containing both info and their respective hash
     * @param key Key to be used for hashing
     * @return true if the Message.hash object is equal to the hash of the message object
     * @throws NoSuchAlgorithmException should never happen, could in case of a library failure
     * @throws InvalidKeyException if an invalid key is provided
     */
    public static boolean verifyHash(Message message, SecretKey key) throws NoSuchAlgorithmException, InvalidKeyException {
        // computedHash is the HMAC of the received plaintext
        byte[] computedHash = hashMessageArr(message, key);
        // receivedHash is the HMAC that was sent by the communication partner
        byte[] receivedHash = decodeB64(message.hash);
        return MessageDigest.isEqual(computedHash, receivedHash);
    }

    /**
     * produces a message hash using HmacSHA256
     * @param message message to be hashed
     * @param key key to initialize hmac with
     * @return byte[] of resulting hash bytes
     * @throws NoSuchAlgorithmException should never happen, could in case of a library failure
     * @throws InvalidKeyException if an invalid key is provided
     */
    private static byte[] hashMessageArr(Message message, SecretKey key) throws NoSuchAlgorithmException, InvalidKeyException {
        // make sure to use the right ALGORITHM for what you want to do (see text)
        Mac hMac = Mac.getInstance("HmacSHA256");
        hMac.init(key);
        // MESSAGE is the message to sign in bytes
        hMac.update(message.allFields().getBytes(StandardCharsets.UTF_8));
        return hMac.doFinal();
    }
}
